#define _CRT_SECURE_NO_WARNINGS
#include <stdio.h>  
#include <stdint.h>  
#include <string.h>  
#include <stdlib.h>  
#include "sm3.h"

// SM3 初始向量  
const uint8_t IV[32] = {
    0x73, 0x80, 0x16, 0x6F, 0x49, 0x14, 0xB2, 0xB9,
    0x17, 0x24, 0x42, 0xD7, 0xDA, 0x8A, 0x06, 0x00,
    0xA9, 0x6F, 0x30, 0xBC, 0x16, 0x31, 0x38, 0xAA,
    0xE3, 0x8D, 0xEE, 0x4D, 0xB0, 0xFB, 0x0E, 0x4E
};

// 循环左移（32位）
uint32_t ROTL(uint32_t x, int n) {
    return (x << n) | (x >> (32 - n));
}

// 常量函数 Tj  
uint32_t Tj(int j) {
    return (j <= 15) ? 0x79CC4519 : 0x7A879D8A;
}

// 布尔函数 FFj  
uint32_t FFj(int j, uint32_t X, uint32_t Y, uint32_t Z) {
    return (j <= 15) ? (X ^ Y ^ Z) : ((X & Y) | (X & Z) | (Y & Z));
}

// 布尔函数 GGj  
uint32_t GGj(int j, uint32_t X, uint32_t Y, uint32_t Z) {
    return (j <= 15) ? (X ^ Y ^ Z) : ((X & Y) | (~X & Z));
}

// 非线性变换函数 P0  
uint32_t P0(uint32_t X) {
    return X ^ ROTL(X, 9) ^ ROTL(X, 17);
}

// 非线性变换函数 P1  
uint32_t P1(uint32_t X) {
    return X ^ ROTL(X, 15) ^ ROTL(X, 23);
}

// 扩展函数 EB  
void EB(const uint8_t Bi[64], uint32_t W[68], uint32_t W1[64]) {
    // 将 Bi 分为 W0~W15  
    for (int i = 0; i < 16; ++i) {
        W[i] = (Bi[i * 4] << 24) | (Bi[i * 4 + 1] << 16) |
            (Bi[i * 4 + 2] << 8) | (Bi[i * 4 + 3]);
    }

    // 扩展 W16~W67  
    for (int j = 16; j <= 67; ++j) {
        W[j] = P1(W[j - 16] ^ W[j - 9] ^ ROTL(W[j - 3], 15)) ^
            ROTL(W[j - 13], 7) ^ W[j - 6];
    }

    // 计算 W1  
    for (int j = 0; j < 64; ++j) {
        W1[j] = W[j] ^ W[j + 4];
    }
}

// 压缩函数 CF  
void CF(const uint8_t Vi[32], const uint8_t Bi[64], uint8_t Vi1[32]) {
    uint32_t W[68] = { 0 };
    uint32_t W1[64] = { 0 };

    EB(Bi, W, W1);

    // 将 Vi 分为 A, B, C, D, E, F, G, H  
    uint32_t R[8];
    for (int i = 0; i < 8; ++i) {
        R[i] = (Vi[i * 4] << 24) | (Vi[i * 4 + 1] << 16) |
            (Vi[i * 4 + 2] << 8) | (Vi[i * 4 + 3]);
    }

    uint32_t A = R[0], B_val = R[1], C = R[2], D = R[3];
    uint32_t E = R[4], F = R[5], G = R[6], H = R[7];
    uint32_t SS1, SS2, TT1, TT2;

    for (int j = 0; j < 64; ++j) {
        SS1 = ROTL((ROTL(A, 12) + E + ROTL(Tj(j), j % 32)), 7);
        SS2 = SS1 ^ ROTL(A, 12);
        TT1 = FFj(j, A, B_val, C) + D + SS2 + W1[j];
        TT2 = GGj(j, E, F, G) + H + SS1 + W[j];
        D = C;
        C = ROTL(B_val, 9);
        B_val = A;
        A = TT1;
        H = G;
        G = ROTL(F, 19);
        F = E;
        E = P0(TT2);
    }

    // 将 ABCDEFGH 重新打包  
    R[0] = A; R[1] = B_val; R[2] = C; R[3] = D;
    R[4] = E; R[5] = F; R[6] = G; R[7] = H;

    uint8_t ABCDEFGH[32];
    for (int i = 0; i < 8; ++i) {
        ABCDEFGH[i * 4] = (R[i] >> 24) & 0xFF;
        ABCDEFGH[i * 4 + 1] = (R[i] >> 16) & 0xFF;
        ABCDEFGH[i * 4 + 2] = (R[i] >> 8) & 0xFF;
        ABCDEFGH[i * 4 + 3] = R[i] & 0xFF;
    }

    // Vi1 = ABCDEFGH ^ Vi  
    for (int i = 0; i < 32; ++i) {
        Vi1[i] = ABCDEFGH[i] ^ Vi[i];
    }
}

// 参数 m 是原始数据，ml 是数据长度（字节数），r 是输出参数，存放 hash 结果  
void SM3Hash(const uint8_t* m, int ml, uint8_t r[32]) {
    uint64_t l = (uint64_t)ml * 8;
    int k = (448 - (l + 1)) % 512;
    if (k < 0) {
        k += 512;
    }

    int total_bits = l + 1 + k + 64;
    int n = total_bits / 512;

    int m1l = n * 512 / 8; // 填充后的长度，512 位的倍数  
    uint8_t* m1 = (uint8_t*)calloc(m1l, sizeof(uint8_t));
    if (m1 == NULL) {
        fprintf(stderr, "Memory allocation failed.\n");
        exit(1);
    }
    memcpy(m1, m, ml);

    m1[ml] = 0x80; // 消息后补 1（10000000）  

    // 添加长度 l 的 64 位大端表示  
    for (int i = 0; i < 8; ++i) {
        m1[m1l - 1 - i] = (l >> (i * 8)) & 0xFF;
    }

    // 将填充后的消息 m′ 按 512 比特进行分组  
    const int BLOCK_SIZE = 64; // 512 位 / 8 = 64 字节  
    uint8_t V[32];
    memcpy(V, IV, 32);

    for (int i = 0; i < n; ++i) {
        CF(V, m1 + i * BLOCK_SIZE, V);
    }

    memcpy(r, V, 32);

    free(m1);
}


// 打印缓冲区  
void dumpbuf(const uint8_t* buf, int len) {
    printf("len=%d\n", len);
    for (int i = 0; i < len; i++) {
        printf("%02x ", buf[i]);
        if ((i + 1) % 16 == 0)
            putchar('\n');
    }
    if (len % 16 != 0)
        putchar('\n');
}
